import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import os

# 额外导入用于画板的库
import tkinter as tk
from PIL import Image, ImageDraw, ImageOps
import numpy as np

# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 定义神经网络模型
class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, kernel_size=5)  # 输出: 6 x 24 x 24
        self.pool1 = nn.AvgPool2d(2)                 # 输出: 6 x 12 x 12
        self.conv2 = nn.Conv2d(6, 16, kernel_size=5) # 输出: 16 x 8 x 8
        self.pool2 = nn.AvgPool2d(2)                 # 输出: 16 x 4 x 4

        self.fc1 = nn.Linear(16 * 4 * 4, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = self.pool1(x)
        x = torch.relu(self.conv2(x))
        x = self.pool2(x)
        x = x.view(-1, 16 * 4 * 4)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        return self.fc3(x)


model = LeNet().to(device)

model_path = "mnist_model.pth"

# 如果模型文件不存在就训练并保存
if not os.path.exists(model_path):
    print("未找到模型，开始训练...")
    transform = transforms.ToTensor()
    train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
    test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)

    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    for epoch in range(30):
        model.train()
        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            loss = criterion(outputs, labels)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        print(f"Epoch {epoch+1}/5, Loss: {loss.item():.4f}")

    torch.save(model.state_dict(), model_path)
    print("模型已保存。")
else:
    print("加载已保存模型...")
    model.load_state_dict(torch.load(model_path, map_location=device))

model.eval()

# 创建画板GUI
class App(tk.Tk):
    def __init__(self):
        super().__init__()
        self.title("手写数字识别")
        self.canvas = tk.Canvas(self, width=280, height=280, bg='white')
        self.canvas.pack()
        self.button = tk.Button(self, text="识别", command=self.predict_digit)
        self.button.pack()
        self.clear_btn = tk.Button(self, text="清除", command=self.clear)
        self.clear_btn.pack()
        self.label = tk.Label(self, text="", font=("Arial", 16))
        self.label.pack()

        self.image = Image.new("L", (280, 280), 255)
        self.draw = ImageDraw.Draw(self.image)

        self.canvas.bind("<B1-Motion>", self.paint)

    def paint(self, event):
        x, y = event.x, event.y
        r = 8
        self.canvas.create_oval(x - r, y - r, x + r, y + r, fill='black')
        self.draw.ellipse([x - r, y - r, x + r, y + r], fill=0)

    def clear(self):
        self.canvas.delete("all")
        self.image = Image.new("L", (280, 280), 255)
        self.draw = ImageDraw.Draw(self.image)
        self.label.config(text="")

    def predict_digit(self):
        img = self.image.resize((28, 28)).convert('L')
        img = ImageOps.invert(img)
        img = np.array(img).astype(np.float32) / 255.0
        img_tensor = torch.tensor(img).unsqueeze(0).unsqueeze(0).to(device)

        with torch.no_grad():
            output = model(img_tensor)
            pred = torch.argmax(output, 1).item()

        self.label.config(text=f"识别结果: {pred}")

if __name__ == "__main__":
    app = App()
    app.mainloop()
